
"""
大模型调用
"""
from langchain_openai import ChatOpenAI
from config import settings
from loguru import logger
from retry import retry


class LLM(object):

    def __init__(self):
        self.llm = ChatOpenAI(
            model=settings.LLM_MODEL,
            # extra_body={"thinking": {"type": "disabled"}}
            extra_body = {"thinking": {"type": "enabled"}}
        )

    @retry(tries=3, delay=1)
    def generate(self, prompt):
        logger.info(f" 提示词： ---> {len(prompt)}")
        result = self.llm.invoke(prompt)
        logger.info(f" 输出： ---> {len(result.content)}")
        return result.content


if __name__ == '__main__':
    llm = LLM()
    print(llm.generate("请给我不是剑神这本小说定义特点标签，最能反映这本小说特点的多个标签，20字以内"))